#!/usr/bin/python

###############################
## A base function for plots ##
###############################

#######################################################      
## Set PyRoot Environment and PlotConfig Path first ###
#######################################################
import sys
ROOTSYS = '/afs/atlas.umich.edu/opt/root/lib'
sys.path.append(ROOTSYS)
PlotConfig='../'
PlotConfig2='.'
sys.path.append(PlotConfig)
sys.path.append(PlotConfig2)

#####################
## Import Module  ###
#####################
import array
import os
import glob
from math import sqrt,fabs,sin
from ROOT import TFile,TTree,TChain,TBranch,TH1,TH1F,TList
from ROOT import TLorentzVector,TGraphAsymmErrors,TMath
from ROOT import THStack,TCanvas,TLegend,TColor,TPaveText,TPad
from ROOT import gStyle,gDirectory
from ROOT import Double

def plot_template_syst(filetag='',hname='',oname='',configfile='plot_configure',hrange='',hxt='',hyt='',hrebin=1,xmax=-999,xmin=-999,ymax=-999,ymin=-999,log=-1,ratio=-1,ratio_min=-999,ratio_max=-999,ratio_adj=1,ratio_err=1,show_mcerr=-1,percent_mcerr=-1,verbose=1,sxt=0.06,syt=0.06,offxt=1.2,offyt=1.2,showATLAS='',**kw):

    '''
    flietag => only include histogram files containing certain tags; 
                in most cases, don't need to set this
                Note that all the histogram should be in ../Result/XXX/
    hname   => the histogram name; to find all the names, simple use 
                "python getlist_histo.py" (loo at histo.list then)
    oname   => the output picture name; if not set, will use the histogram name
    configfile => Which config file to use, default value is "plot_configure"
    hrange  => A string indicating the ORIGINAL histogram ranges: "nbins,xmin,xmax"; 
                e.g.  "1000,-10,10"
               If you want the program to read the histogram range automatically from input histo, 
               do not set this parameter
    hxt     => The title of x axis
    hyt     => The title of y axis    
    hrebin  => How many times for histogram rebinning
    xmin    => the lower limit for x axis
    xmax    => the upper limit for x axis
    ymax    => The maximum of y axis; if set to -1, will use 20*ActualMaximun for logrithm plot 
                and 2*ActualMaximum for linear plot
    ymin    => The minimum of y axis    
    log     => By default, output plots will be both log and linear; if you want only log plots, set it to be 1, else set it to be 0
    ratio   => If you want to also have ratio plot in the canvas, set it to a non-zero value; 
                in that case, Top part of the picture will show Data VS MC, and bottom part 
                wil show the ratio of Data over MC(total). The fraction of bottom plot over top plot 
                is defined as this "ratio" variable
    show_mcerr => If you want to show MC uncertainties with an error band, set this variable to be 1
    percent_mcerr => Used together with "show_mcerr" to set the fractional total uncertainty of MC prediction
    sxt,syt,offxt,offyt => x axis title size / y axis title size / x axis title offset / y axis title offset
    showATLAS  => If you want to show "ATLAS"/"ATLAS Preliminary" here, put the actual string here; otherwise leave it blank
    verbose => Set it to be 1, if you want it output verbose message
    '''

    print 'INFO => Make Plot for histogram %s' % hname
    print ''

    
    ############################
    ## Load Atlas Plot Style ###            
    ############################
    
    from plot_style import atlas_style
    atlas_style()

    ###################
    ## Define Canvas ##
    ###################
    print 'INFO => Defining Canvas'
    
    # If want ratio plot to be shown at the bottom, define its fraction first
    if ratio<0: canvas_w,canvas_h=800,600
    elif ratio>0: canvas_w,canvas_h=800,int(600*(1+ratio))
    else: 
        print 'ERROR => The scale of RATIO plot is wrongly defined: %.4f' % ratio 
        sys.exit(-1)
    # Define canvas and create sub pad if needed
    gStyle.SetPadLeftMargin(0.15)
    gStyle.SetPadRightMargin(0.10)
    gStyle.SetPadTopMargin(0.05)
    gStyle.SetPadBottomMargin(0.15)
    MyC=TCanvas('MyC','MyC',canvas_w,canvas_h)
    MyC.SetTicks(1,1)    
    if ratio>0:
        fraction=ratio+0.2
        Pad1 = TPad("p1","p1",0,fraction*1.0/(fraction+1),1,1,0,0) # x1,y1,x2,y2
        Pad1.SetMargin(0.15,0.10,0.03,0.05)
        Pad2 = TPad("p2","p2",0,0,1,fraction*1.0/(fraction+1),0,0)
        Pad2.SetMargin(0.15,0.10,0.15/fraction,0)
        Pad1.Draw()
        Pad2.Draw()
        Pad2.SetGrid()
        print 'INFO => Canvas is defined as (width=%d,height=%d): Ratio Plot Pad / Main Pad = %.3f' % (canvas_w,canvas_h,ratio)
    else:
        print 'INFO => Canvas is defined as (width=%d,height=%d)' % (canvas_w,canvas_h)
    
    ##########################################
    ## Load files,Inputs and Get histograms ##
    ##########################################
    print 'INFO => Loading histograms'
    
    # Parse config file
    from plot_config import plot_config
    dict=plot_config(configfile)
    # Read variables
    IntLumi=dict['IntLumi'] # Integrated Luminosity,fb-1
    MCScale=dict['MCScale'] # Global MCScale Factor
    Dataset=dict['Datasets'] # The Dataset contribution, data/sig/zz/wz...
    
    # Get the histogram nbins, xmin and xmax; And declasre histograms for further usage
    histo_range=hrange.split(',')
    histo_all={}
    histo_key={}
    histo_nbin, histo_xmin, histo_xmax = 0, 0., 0.
    if len(histo_range)==3: # if hrange is specified
        histo_nbin, histo_xmin, histo_xmax = int(histo_range[0]),float(histo_range[1]),float(histo_range[2])
    else: # if hrange is not specified, grab a histogram file to extract the histogram range info
        for dataset in Dataset:
            for key in dict[dataset].keys():
                if key=='Color' or key=='Corr' or key=='OverLay': continue
                # get list of folders for this dataset or process
                list_folders=[]
                if 'data' in key: list_folders.append(key)
                else:
                    for folder in dict[dataset][key].keys(): list_folders.append(folder)
                for folder in list_folders:
                    # Read histo file
                    if filetag: filename=glob.glob('%s/*%s*root' % (folder,filetag))
                    else: filename=glob.glob('%s/*root' % (folder))
                    file=TFile(filename[0])
                    histo_temp=gDirectory.Get(hname)
                    if not histo_temp:
                        print 'ERROR ==> In %s, no histogram %s exist!' % (filename,hname)
                        continue
                    histo_nbin=histo_temp.GetNbinsX()
                    histo_xmin=histo_temp.GetXaxis().GetXmin()
                    histo_xmax=histo_temp.GetXaxis().GetXmax()
                    file.Close()
                    break
                break
            break    
    for dataset in Dataset:
        histo_all[dataset]=TH1F(dataset,dataset,histo_nbin, histo_xmin, histo_xmax)
        histo_all[dataset].Sumw2()
    for dataset in Dataset:
        histo_key[dataset]={}
        for key in dict[dataset].keys():
            if key=='Color' or key=='Corr' or key=='OverLay': continue
            histo_key[dataset][key]=TH1F('%s_%s' % (dataset,key),'%s_%s' % (dataset,key),histo_nbin, histo_xmin, histo_xmax)
            histo_key[dataset][key].Sumw2()
    # loop all files and merge histograms (normalization will be done for MC)
    for dataset in Dataset:
        print 'INFO ==> Looking at %s' % dataset
        this_corr=dict[dataset]['Corr']
        for key in dict[dataset].keys():
            if key=='Color' or key=='Corr' or key=='OverLay': continue
            # get list of folders for this dataset or process
            if verbose!=-1:  print 'INFO ==> KEY = %s' % key
            list_folders=[]
            if 'data' in key: list_folders.append(key)
            else: 
                for folder in dict[dataset][key].keys(): list_folders.append(folder)
            # loop all the folders one by one and merge histograms
            xs, k, fe, entry = 0., 0., 0., 0.
            for folder in list_folders:
                # Read histo file
                if filetag: filename=glob.glob('%s/*%s*root' % (folder,filetag))
                else: filename=glob.glob('%s/*root' % (folder))
                file=TFile(filename[0])        
                # Get histogram
                histo_temp=gDirectory.Get(hname)
                if not histo_temp: 
                    print 'ERROR ==> In %s, no histogram %s exist!' % (filename,hname)
                    continue
                # For MC, mormalize histogram and then combine
                if 'mc' in folder:
                    xs=float(dict[dataset][key][folder]['Xsection(pb)'])
                    k=float(dict[dataset][key][folder]['k-factor'])
                    fe=float(dict[dataset][key][folder]['FilterEff(gen)'])
                    entry+=float(dict[dataset][key][folder]['Nevts(Weight)'])
                    if verbose!=-1: print 'INFO ==> Add histograms from %s' % folder
                    histo_key[dataset][key].Add(histo_temp)
                # For Data, just combine histogram inside one dataset
                if 'data' in key:
                    if verbose!=-1: print 'INFO ==> Add histograms from %s' % folder
                    histo_all[dataset].Add(histo_temp)
	            # Close histo file
        	    file.Close()
            # If MC contains several different subsamples in one MCID, one needs to sum 
            # the entries to normalize MC histograms, scale factor = 
            # xs*kfactor*fe*Lumi*MCScale*dict[dataset][Corr]/entry
            if not 'data' in key:
                mc_scale=xs*k*fe*IntLumi*1000.0*MCScale*this_corr/entry
                histo_key[dataset][key].Scale(mc_scale)
                histo_all[dataset].Add(histo_key[dataset][key])
                if verbose!=-1: print 'INFO ==> Normalization done for %s -> %.6f' % (key, mc_scale)

    ####################
    ## Begin to plot ###
    ####################
    from plot_style import color_convert

    # Set Histogram Fill style
    gStyle.SetHistFillStyle(1001)

    # Set Attributes for histograms:
    # 1) Rebin
    # 2) Set Colors
    # 3) get the maximum entry for latter usage
    # 4) stack histograms
    # 5) sum up all MCs into one, prepare for ratio plot
    hs=THStack("hs","Stacked 1D histograms")
    histo_allMC=TH1F('MCALL','MCALL',histo_nbin, histo_xmin, histo_xmax)
    histo_allMC.Rebin(hrebin)    
    histo_allMC.Sumw2()

    binwidth=(histo_xmax-histo_xmin)/histo_nbin
    binwidth*=hrebin
    if xmin!=-999:    hminb=int((xmin-histo_xmin)/binwidth)+1
    else: hminb=1
    if xmax!=-999:    hmaxb=int((xmax-histo_xmin)/binwidth)
    else:  hmaxb=histo_nbin
    
    max=0.
    for dataset in Dataset:
        # Rebin
    	histo_all[dataset].Rebin(hrebin)
    	# Set Color
    	color_str=dict[dataset]['Color']
    	color=color_convert(color_str)
    	if color==-1: 
      		print 'ERROR ==> Can not parse the color: %s' % dict[dataset]['Color']
      		sys.exit(1)
    	print 'INFO ==> Set Color %s for dataset:%s' % (color_str,dataset)
        if dict[dataset]['OverLay']=='0':
            histo_all[dataset].SetFillColor(color)
        else:
            histo_all[dataset].SetLineColor(color)
    	# Add into THStack
        if dataset!='Data' and dict[dataset]['OverLay']=='0':
            hs.Add(histo_all[dataset])
            histo_allMC.Add(histo_all[dataset])
    # Make a ratio plot
    histo_ratio=histo_all['Data'].Clone()
    histo_ratio.SetName('RATIO')
    histo_ratio.Divide(histo_allMC)  
    # Get yaxis max
    for i in range(hminb,hmaxb+1):
        if max<histo_allMC.GetBinContent(i):
            max=histo_allMC.GetBinContent(i)
        if max<histo_all['Data'].GetBinContent(i):
            max=histo_all['Data'].GetBinContent(i)
            
    # Syst. uncertainty box
    NbinsX=histo_ratio.GetNbinsX()
    histo_ratioErr=histo_ratio.Clone()
    for i in range(NbinsX):
        histo_ratioErr.SetBinContent(i+1, 1.)  #meanvalue = 1
        histo_ratioErr.SetBinError(i+1, 0.039) #luminosity error: 3.9%

    ###############
    ## Plotting ###
    ###############
    print  'INFO => Now make plots'
    
    # Set Plot attibute for stacked histogram and Draw the MC part
    if len(Dataset)==1 and Dataset[0]=="Data":
        pass
    else:         
        if ratio>0: Pad1.cd()
        else: MyC.cd()
        hs.Draw('HIST')
        if show_mcerr==1:
            if percent_mcerr>0:
              for i in range(histo_all['Data'].GetNbinsX()):
                  histo_allMC.SetBinError(i+1,percent_mcerr*histo_allMC.GetBinContent(i+1))
            histo_allMC.SetFillStyle(3357)
            gStyle.SetHatchesLineWidth(2)
            histo_allMC.SetFillColor(15)
            histo_allMC.SetMarkerSize(0)
            histo_allMC.Draw('E2SAME')
        for dataset in Dataset:
            if dict[dataset]['OverLay'] == '1':
                 print dataset
                 histo_all[dataset].Draw('HISTSAME')
        if ratio_err>0:
            Pad2.cd()
            histo_ratioErr.SetMarkerSize(0.)
            histo_ratioErr.SetFillColor(5)
            histo_ratioErr.Draw("LE2")
        if ratio>0: 
            Pad2.cd()
            histo_ratio.SetMarkerSize(0.8)
            histo_ratio.Draw("LPESAME")
            
    # Set drawing attributes for histograms
    # Set title and its size and offset (handle special case if ratio plot is requested as well)
    hs.GetXaxis().SetRange(hminb,hmaxb)            
    hs.GetYaxis().SetTitleSize(syt)
    hs.GetYaxis().SetTitle("%s" % hyt)
    hs.GetYaxis().SetTitleOffset(offyt)
    if ratio<=0:
        hs.GetXaxis().SetTitleSize(sxt)
        hs.GetXaxis().SetTitle("%s" % hxt)
        hs.GetXaxis().SetTitleOffset(offxt)        
    # Set drawing attributes for RATIO histogram, need to ajust title and label size to be consistent 
    # with the primary plot
    if ratio>0:
        hs.GetXaxis().SetLabelSize(0)
        histo_ratio.GetXaxis().SetRange(hminb,hmaxb) 
        histo_ratio.GetXaxis().SetTitle(hxt)
        histo_ratio.GetXaxis().SetTitleSize(sxt*1.0/fraction)
        histo_ratio.GetXaxis().SetTitleOffset(offxt)
        histo_ratio.GetYaxis().SetTitle('data/MC')
        histo_ratio.GetYaxis().SetTitleSize(syt*1.0/fraction)
        histo_ratio.GetYaxis().SetTitleOffset(offyt*fraction)    
        histo_ratio.GetXaxis().SetLabelSize(histo_all['Data'].GetLabelSize()*1.0/fraction)
        histo_ratio.GetYaxis().SetLabelSize(histo_all['Data'].GetLabelSize()*1.0/fraction)
        rmin,rmax=10e9,0
        for i in range(hminb,hmaxb+1):
          if histo_ratio.GetBinContent(i)==0: continue
          if rmin>histo_ratio.GetBinContent(i):
              rmin=histo_ratio.GetBinContent(i)-histo_ratio.GetBinError(i)
          if rmax<histo_ratio.GetBinContent(i):
              rmax=histo_ratio.GetBinContent(i)+histo_ratio.GetBinError(i)
        if rmin<0:   rmin=0
        if rmax>2:   rmax=2
        if ratio_adj==1:
          if ratio_min!=-999 and rmin<ratio_min:  rmin=ratio_min
          if ratio_max!=-999 and rmax>ratio_max:  rmax=ratio_max
        else:
          if ratio_min!=-999:  rmin=ratio_min
          if ratio_max!=-999:  rmax=ratio_max
        histo_ratio.SetMinimum(rmin)
        histo_ratio.SetMaximum(rmax)   
        histo_ratioErr.GetXaxis().SetTitleSize(sxt*1.0/fraction)
        histo_ratioErr.GetXaxis().SetTitleOffset(offxt)
        histo_ratioErr.GetYaxis().SetTitleSize(syt*1.0/fraction)
        histo_ratioErr.GetYaxis().SetTitleOffset(offyt*fraction)    
        histo_ratioErr.GetXaxis().SetLabelSize(histo_all['Data'].GetLabelSize()*1.0/fraction)
        histo_ratioErr.GetYaxis().SetLabelSize(histo_all['Data'].GetLabelSize()*1.0/fraction)
        histo_ratioErr.SetMinimum(rmin)
        histo_ratioErr.SetMaximum(rmax)   
        histo_ratioErr.GetXaxis().SetRange(hminb,hmaxb) 

        # automatically adjust y axis stick intensity
        histo_ratio.GetYaxis().SetNdivisions(5+100*5)
        histo_ratioErr.GetYaxis().SetNdivisions(5+100*5)
        
    MyC.Update()
               
    ###############
    ## Add Text ###
    ###############
        
    # Draw Text (Lumi)(Part1)
    if ratio>0: Pad1.cd()
    else: MyC.cd()    
    tl = TPaveText(0.17,0.71,0.18,0.81,"NDC")
    tl.SetBorderSize(0)
    tl.SetFillColor(10)
    tl.SetTextColor(1)
    tl.SetTextSize(0.03)
    tl.SetTextAlign(22)
    tl.AddText('#int')
    tl.Draw()
    # Draw Text (Lumi)(Part2)
    t2 = TPaveText(0.18,0.72,0.41,0.82,"NDC")
    t2.SetBorderSize(0)
    t2.SetFillColor(10)
    t2.SetTextColor(1)
    t2.SetTextSize(0.05)
    t2.SetTextAlign(12)
    t2.AddText('Ldt = %.2ffb^{-1}' % IntLumi)
    t2.Draw()
    # Draw Text (ATLAS )
    ta = TPaveText(0.17,0.82,0.41,0.92,"NDC")
    ta.SetBorderSize(0)
    ta.SetFillColor(10)
    ta.SetTextColor(1)
    ta.SetTextSize(0.05)
    ta.SetTextAlign(12)
    ta.SetTextFont(72)
    ta.AddText(showATLAS)
    ta.Draw()
    # Draw Text (ATLAS)
    ta1 = TPaveText(0.38,0.72,0.55,0.82,"NDC")
    ta1.SetBorderSize(0)
    ta1.SetFillColor(10)
    ta1.SetTextColor(1)
    ta1.SetTextSize(0.05)
    ta1.SetTextAlign(12)
    ta1.AddText('#sqrt{s}=7TeV')
    ta1.Draw()     
    # Draw TLegend
    # how high each item is in the legend
    bhl=0.045
    shl=0.04
    # text size 
    bts=bhl
    sts=shl
    # how many items in each column if the legend has to be split into two columns
    ncol=6
    # Use number of entries to determine the y-size of legend automatically
    if ratio>0: legend_scale=1.0/(1+ratio)
    else: legend_scale=1
    ndtsets=len(Dataset)
    nentries=ndtsets
    if show_mcerr==1: nentries = nentries+1
    if nentries<=7:
      lg = TLegend(0.60, 0.92-bhl*nentries, 0.85, 0.92)
      lg.SetBorderSize(0)
      lg.SetFillColor(10)
      lg.SetTextSize(bts)
      for i in range(ndtsets):
          dataset=''
          if i==0:
              dataset=Dataset[0]
              lg.AddEntry(histo_all[dataset],dataset,"LPE") #Data
          else:
              dataset=Dataset[ndtsets-i]
              if dict[dataset]['OverLay']=='1':
                   lg.AddEntry(histo_all[dataset],dataset,'l')
              else:
                   lg.AddEntry(histo_all[dataset],dataset,"f")
      if show_mcerr==1:
        if percent_mcerr>0:  lg.AddEntry(histo_allMC, "#sigma_{stat+syst}", "f")
        else:                lg.AddEntry(histo_allMC, "#sigma_{stat}", "f")
      lg.Draw()
    else:
      lg1 = TLegend(0.52, 0.92-shl*ncol, 0.70, 0.92)
      lg1.SetBorderSize(0)
      lg1.SetFillColor(10)
      lg1.SetTextSize(sts)
      lg2 = TLegend(0.70, 0.92-shl*(nentries-ncol), 0.88, 0.92)
      lg2.SetBorderSize(0)
      lg2.SetFillColor(10)
      lg2.SetTextSize(sts)
      for i in range(ncol):
          dataset=''
          if i==0:
              dataset=Dataset[0]
              lg1.AddEntry(histo_all[dataset],dataset,"LPE") #Data
          elif i<ncol:
              dataset=Dataset[ndtsets-i]
              if dict[dataset]['OverLay']=='1':
                   lg1.AddEntry(histo_all[dataset],dataset,'l')
              else:
                   lg1.AddEntry(histo_all[dataset],dataset,"f")
      for i in range(ncol,ndtsets):
          dataset=Dataset[ndtsets-i]
          if dict[dataset]['OverLay']=='1':
               lg2.AddEntry(histo_all[dataset],dataset,'l')
          else:
               lg2.AddEntry(histo_all[dataset],dataset,"f")
      if show_mcerr==1:
        if percent_mcerr>0:  lg2.AddEntry(histo_allMC, "#sigma_{stat+syst}", "f")
        else:                lg2.AddEntry(histo_allMC, "#sigma_{stat}", "f")
      lg1.Draw()
      lg2.Draw()

    ##################
    ## Print Plots ###
    ##################

    # logrithmic plot
    if log==1 or log==-1: 
        # set canvas to be loglike
        if ratio>0:
            Pad1.SetLogy(1)
            Pad1.cd()
        else:
            MyC.SetLogy(1)
            MyC.cd()
        # set y axis range
        hs.SetMinimum(1e-1)
        hs.SetMaximum(max*sqrt(hs.GetMaximum()/1e-1)) 
        if ymax!=-999: hs.SetMaximum(ymax)
        if ymin!=-999: hs.SetMinimum(ymin)
        # re-compute errors for data, and draw it
        # Add poission error to data points, need to set assymetry errors to graph
        gerr = TGraphAsymmErrors(histo_all['Data'])
        for i in range(gerr.GetN()):
            var_x, var_y = Double(0.), Double(0.)
            gerr.GetPoint(i,var_x,var_y)
            if var_y == 0.:
                gerr.SetPoint(i,var_x,-1.0)
                gerr.SetPointEYlow(i,-1)
                gerr.SetPointEYhigh(i,-1)
                continue
            err_low = var_y - (0.5*TMath.ChisquareQuantile(0.1586555,2.*var_y))
            err_high = (0.5*TMath.ChisquareQuantile(1.-0.1586555,2.*(var_y+1))) - var_y
            gerr.SetPointEYlow(i, err_low)
            gerr.SetPointEYhigh(i, err_high)
            #print "Point=%d, Value=%d, err_low=%.2f, err_high=%.2f\n" % (i, var_y, err_low, err_high)
        gerr.SetMarkerColor(1)
        gerr.SetMarkerSize(1)
        gerr.SetMarkerStyle(20)
        gerr.Draw('PE1SAME')
        MyC.Update()
        # print the canvas
        if oname: plot_name='%s_log.png' % oname
        else: plot_name='%s_log.png' % hname                 
        print 'INFO ==> Print canvas to file: %s' % plot_name
        MyC.Print(plot_name)
    # linear plot
    if log==0 or log==-1:
        # set canvas to be loglike
        if ratio>0:
            Pad1.SetLogy(0)
            Pad1.cd()
        else:
            MyC.SetLogy(0)
            MyC.cd()
        # set y axis range
        hs.SetMaximum(max*2)
        if ymax!=-999: hs.SetMaximum(ymax)
        if ymin!=-999: hs.SetMinimum(ymin)
        # draw data
        histo_all['Data'].SetMarkerColor(1)
        histo_all['Data'].SetMarkerStyle(20)    
        histo_all['Data'].Draw('LPE1SAME')
        MyC.Update()
        # print the canvas
        if oname: plot_name='%s_linear.png' % oname
        else: plot_name='%s_linear.png' % hname                 
        print 'INFO ==> Print canvas to file: %s' % plot_name
        MyC.Print(plot_name)

    #####################################################
    ## Clear the histograms and canvas after plotting ###
    #####################################################
  
    for dataset in histo_all.keys():
        histo_all[dataset].SetDirectory(0)
        TH1.AddDirectory(0)
    MyC.Close()
